#
# common for CIFAR10
#
###############################################################################
#
import os
import torchvision
from typing import Any, Callable, Dict, List, Optional, Tuple

# import torch.nn as nn
# import torch.nn.functional as F
# Fu.relu()是函数调用，一般使用在foreward函数里。
# nn.ReLU()是模块调用，一般在定义网络层的时候使用。


ROOT = f"e:\sfxData\DeepLearning\cifar-10"
CIFAR10_TRANSFORM_DEFAULT = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),                      # 转换为张量
    # torchvision.transforms.Normalize((0.1307,), (0.3081,))  # 设定标准化值
])


def get_dataset(root: str, train: bool = True, transform: Optional[Callable] = None) -> torchvision.datasets.CIFAR10:
    """
    取得 CIFAR10 数据集

    参数
    :root - 数据集根目录
    :train - 是否取得训练集
    :transform - 转换
    """
    ds = __cifar10(root, train=train, transform=transform)
    return ds


def __cifar10(root: str, train: bool = True, transform: Optional[Callable] = None) -> torchvision.datasets.CIFAR10:
    """
    取得 CIFAR10 数据集.

    参数
    :root - 数据集根目录
    :train - 是否取得训练集
    :transform - 转换
    """
    xtransform = CIFAR10_TRANSFORM_DEFAULT if transform == None else transform
    isDownload = not(os.path.exists(root)) or not os.listdir(root)
    dset = torchvision.datasets.CIFAR10(root=root, train=train,
                                        transform=xtransform, download=isDownload)
    return dset
